import numpy as np
from predict_config import get_config
from vgg13_shortcuts import vgg13_shortcuts_v2
import os
from easy_io import write_pkl_file
import global_config
from keras.backend.tensorflow_backend import set_session
import tensorflow as tf


def predict_model(candidate_file, candidate_vol_file, batchsize, crop_shape, gpus,
                  weight_decay, prob_thresh, model_name, model_epoch, train_or_test, **kwargs):
    args = locals()
    args.update(args.pop('kwargs'))

    os.environ["CUDA_VISIBLE_DEVICES"] = gpus
    config = tf.ConfigProto()
    config.gpu_options.allow_growth = True
    set_session(tf.Session(config=config))

    print('')
    print('*' * 20)
    for k in sorted(args.keys()):
        print('{k}: {v}'.format(k=k, v=args[k]))
    print('*' * 20)
    print('')

    num_classes = (2,)
    feature_names = ('Mal_2',)
    model = vgg13_shortcuts_v2((64, 64, 64, 1), 32, weight_decay, num_classes, feature_names)

    if train_or_test == 'train':

        for valid_fold in range(4):
            include_testset = 'with_test' in model_name
            config_dict = get_config(candidate_file, candidate_vol_file, batchsize, crop_shape, [valid_fold],
                                     include_testset, prob_thresh)
            valid_generator = config_dict['valid_generator']
            nb_valid_samples = config_dict['nb_valid_samples']
            valid_candidates = config_dict['valid_candidates']
            kwargs['input_shape'] = config_dict['input_shape']

            model_path = os.path.join(global_config.model_folder, model_name+'_fold'+str(valid_fold))
            model_weights = os.listdir(model_path)
            model_weight = [x for x in model_weights if x.startswith(model_epoch)]
            if len(model_weight) == 0:
                print('no model weights!')
            elif len(model_weight) > 1:
                print('too many weights!')
            else:
                model_weight = os.path.join(global_config.model_folder, model_name+'_fold'+str(valid_fold), model_weight[0])
            model.load_weights(model_weight, by_name=True)

            os.makedirs(os.path.join(global_config.result_folder, model_name+'_'+model_epoch), exist_ok=True)
            result_saveto = os.path.join(global_config.result_folder, model_name+'_'+model_epoch, model_name+'_fold'+str(valid_fold)+'.pkl')
            try:
                probs = []
                nb_seen_samples = 0
                while nb_seen_samples < nb_valid_samples:
                    samples = next(valid_generator)
                    probs.append((model.predict_on_batch(samples)))
                    nb_seen_samples += len(samples)
                assert nb_seen_samples == nb_valid_samples
                probs = np.concatenate(probs, axis=0)
                for i in range(len(probs)):
                    valid_candidates[i]['mal_prob'] = probs[i, 1]
                write_pkl_file(result_saveto, valid_candidates)
            except KeyboardInterrupt:
                pass

        if 'with_test' in model_name:
            print("don't need stage1 test set results!")
            return

        for valid_fold in range(4):
            include_testset = 'with_test' in model_name
            config_dict = get_config(candidate_file, candidate_vol_file, batchsize, crop_shape, None,
                                     True, prob_thresh, only_testset=True)
            valid_generator = config_dict['valid_generator']
            nb_valid_samples = config_dict['nb_valid_samples']
            valid_candidates = config_dict['valid_candidates']
            kwargs['input_shape'] = config_dict['input_shape']
            model_path = os.path.join(global_config.model_folder, model_name+'_fold'+str(valid_fold))
            model_weights = os.listdir(model_path)
            model_weight = [x for x in model_weights if x.startswith(model_epoch)]
            if len(model_weight) == 0:
                print('no model weights!')
            elif len(model_weight) > 1:
                print('too many weights!')
            else:
                model_weight = os.path.join(global_config.model_folder, model_name+'_fold'+str(valid_fold), model_weight[0])
            model.load_weights(model_weight, by_name=True)

            os.makedirs(os.path.join(global_config.result_folder, model_name+'_'+model_epoch), exist_ok=True)
            result_saveto = os.path.join(global_config.result_folder, model_name+'_'+model_epoch, model_name+'_merge'+str(valid_fold)+'.pkl')
            try:
                probs = []
                nb_seen_samples = 0
                while nb_seen_samples < nb_valid_samples:
                    samples = next(valid_generator)
                    probs.append((model.predict_on_batch(samples)))
                    nb_seen_samples += len(samples)
                assert nb_seen_samples == nb_valid_samples
                probs = np.concatenate(probs, axis=0)
                for i in range(len(probs)):
                    valid_candidates[i]['mal_prob'] = probs[i, 1]
                write_pkl_file(result_saveto, valid_candidates)
            except KeyboardInterrupt:
                pass

    elif train_or_test == 'test':
        for valid_fold in range(4):
            config_dict = get_config(candidate_file, candidate_vol_file, batchsize, crop_shape, None,
                                     True, prob_thresh)
            valid_generator = config_dict['valid_generator']
            nb_valid_samples = config_dict['nb_valid_samples']
            valid_candidates = config_dict['valid_candidates']
            kwargs['input_shape'] = config_dict['input_shape']

            model_path = os.path.join(global_config.model_folder, model_name+'_fold'+str(valid_fold))
            model_weights = os.listdir(model_path)
            model_weight = [x for x in model_weights if x.startswith(model_epoch)]
            if len(model_weight) == 0:
                print('no model weights!')
            elif len(model_weight) > 1:
                print('too many weights!')
            else:
                model_weight = os.path.join(global_config.model_folder, model_name+'_fold'+str(valid_fold), model_weight[0])
            model.load_weights(model_weight, by_name=True)

            os.makedirs(os.path.join(global_config.result_folder, model_name+'_'+model_epoch+'_stage2'), exist_ok=True)
            result_saveto = os.path.join(global_config.result_folder, model_name+'_'+model_epoch+'_stage2', model_name+'_merge'+str(valid_fold)+'.pkl')
            try:
                probs = []
                nb_seen_samples = 0
                while nb_seen_samples < nb_valid_samples:
                    samples = next(valid_generator)
                    probs.append((model.predict_on_batch(samples)))
                    nb_seen_samples += len(samples)
                assert nb_seen_samples == nb_valid_samples
                probs = np.concatenate(probs, axis=0)
                for i in range(len(probs)):
                    valid_candidates[i]['mal_prob'] = probs[i, 1]
                write_pkl_file(result_saveto, valid_candidates)
            except KeyboardInterrupt:
                pass
